import matplotlib.pyplot as plt
import numpy as np
import pandas as pd
import seaborn as sns

# 设置中文显示
plt.rcParams['font.sans-serif'] = ['Microsoft YaHei']  # 或 'Heiti SC', 'Hiragino Sans GB'
plt.rcParams['axes.unicode_minus'] = False

# 设置样式
# sns.set_style('whitegrid')  # 替换 plt.style.use('seaborn')
# sns.set_palette("husl")

# 消融实验结果数据
ablation_data = {
    "优化组合": ["基准版本", "分块矩阵乘法", "融合核函数", "内存池", "hipBLAS", 
              "分块矩阵乘法+融合核函数", "分块矩阵乘法+内存池", "融合核函数+内存池",
              "hipBLAS+内存池", "所有优化（不包括多DCU）", "所有优化"],
    "平均时间(ms)": [0.255267, 0.250188, 0.216203, 0.142662, 0.251382,
                  0.217472, 0.13776, 0.0929616, 0.123059, 0.0937376, 0.0942341],
    "加速比": [1, 1.0203, 1.18068, 1.78932, 1.01546, 1.17379, 1.85298, 
             2.74594, 2.07436, 2.72321, 2.70886]
}

# 批次大小实验数据
batch_size_data = {
    "批次大小": [32, 64, 128, 256, 512, 1024, 2048, 4096],
    "平均时间(ms)": [0.0940751, 0.093723, 0.093387, 0.0942121, 
                  0.0957505, 0.0937791, 0.0988716, 0.110802]
}

# 隐藏层维度实验数据
hidden_dim_data = {
    "隐藏层维度": [10, 20, 50, 100, 200, 500],
    "平均时间(ms)": [0.090706, 0.0939135, 0.108932, 0.131848, 0.1808, 0.324196]
}

# 转换为DataFrame
df_ablation = pd.DataFrame(ablation_data)
df_batch = pd.DataFrame(batch_size_data)
df_hidden = pd.DataFrame(hidden_dim_data)

# 1. 消融实验 - 加速比
plt.figure(figsize=(8, 6))
bars = plt.barh(df_ablation["优化组合"], df_ablation["加速比"], color='skyblue')
plt.bar_label(bars, fmt='%.2fx', padding=3)
plt.title('不同优化组合的加速比 (相对于基准版本)')
plt.xlabel('加速比')
plt.grid(True, linestyle='--', alpha=0.6)
plt.tight_layout()
plt.savefig('不同优化组合的加速比.png', dpi=300)
plt.close()

# 2. 消融实验 - 执行时间
plt.figure(figsize=(8, 6))
bars = plt.barh(df_ablation["优化组合"], df_ablation["平均时间(ms)"], color='lightgreen')
plt.bar_label(bars, fmt='%.3f', padding=3)
plt.title('不同优化组合的执行时间')
plt.xlabel('时间 (ms)')
plt.grid(True, linestyle='--', alpha=0.6)
plt.tight_layout()
plt.savefig('不同优化组合的执行时间.png', dpi=300)
plt.close()

# 3. 批次大小实验
plt.figure(figsize=(8, 6))
plt.plot(df_batch["批次大小"], df_batch["平均时间(ms)"], marker='o', linestyle='-', color='orange')
plt.title('不同批次大小下的执行时间')
plt.xlabel('批次大小')
plt.ylabel('时间 (ms)')
plt.xscale('log', base=2)
plt.xticks(df_batch["批次大小"], df_batch["批次大小"])
plt.grid(True, which="both", linestyle='--', alpha=0.6)
plt.tight_layout()
plt.savefig('不同批次大小下的执行时间.png', dpi=300)
plt.close()

# 4. 隐藏层维度实验
plt.figure(figsize=(8, 6))
plt.plot(df_hidden["隐藏层维度"], df_hidden["平均时间(ms)"], marker='s', linestyle='-', color='purple')
plt.title('不同隐藏层维度下的执行时间')
plt.xlabel('隐藏层维度')
plt.ylabel('时间 (ms)')
plt.xscale('log')
plt.xticks(df_hidden["隐藏层维度"], df_hidden["隐藏层维度"])
plt.grid(True, which="both", linestyle='--', alpha=0.6)
plt.tight_layout()
plt.savefig('不同隐藏层维度下的执行时间.png', dpi=300)
plt.close()